#!/opt/local/bin/python
from PyQt4.QtCore import *
from PyQt4.QtGui import *
import dependencies as dep
import pprint
from os.path import exists, expanduser, islink
from sys import path, argv
from os import chdir, getcwd, rename, listdir, remove, link
from devices import devices
chdir(getcwd() + '/iphone-dataprotection')
path.append(getcwd() + '/python_scripts')
from time import sleep
from hashlib import sha1
from re import search
from subprocess import call, check_output, check_call
from urllib import urlretrieve
import ios_examiner
import kernel_patcher

global IPSW_path 
global ramdiskMeth
global IPSW_files


class MainWindow(QMainWindow):

	def __init__(self):

		self.iThing = 'none_set'

		super(MainWindow, self).__init__(parent = None)
#Message Box

		#logDockWidget = QDockWidget("Log", self)
		#logDockWidget.setFeatures(logDockWidget.DockWidgetMovable | logDockWidget.DockWidgetMovable | logDockWidget.DockWidgetFloatable)
		#logDockWidget.setObjectName("LogDockWidget")
		#logDockWidget.setAllowedAreas(Qt.LeftDockWidgetArea|Qt.RightDockWidgetArea)


		mainWidget = QWidget()
		mainHBox = QHBoxLayout()
		self.listWidget = QListWidget()
		self.listWidget.setFixedWidth(400)

		mainHBox.addWidget(self.listWidget)
		self.stackedWidget = QStackedWidget()


		vboxSetup = QVBoxLayout()
		setupButton = QPushButton("Setup")
		continueButton = QPushButton("Continue")
		self.connect(setupButton, SIGNAL("clicked()"), self.setup)
		self.connect(continueButton, SIGNAL("clicked()"), self.displayiThing)
		vboxSetup.addWidget(setupButton)
		vboxSetup.addWidget(continueButton)
		self.groupBoxStart = QGroupBox("Start")
		self.groupBoxStart.setLayout(vboxSetup)
		self.stackedWidget.addWidget(self.groupBoxStart)
			
		#Get selected radio button by isselected
		vboxiThing = QVBoxLayout()
		self.threeGradio = QRadioButton("iPhone 3G")
		self.threeGSradio = QRadioButton("iPhone 3GS")
		self.fourGSMradio = QRadioButton("iPhone 4 GSM AT&T")
		self.fourCDMAradio = QRadioButton("iPhone 4 CDMA Verizon")
		self.pod2Gradio = QRadioButton("iPod 2G")
		self.podthreeGradio = QRadioButton("iPod 3G")
		self.pod4Gradio = QRadioButton("iPod 4G")
		self.pad1radio = QRadioButton("iPad 1")
		self.skip1 = QRadioButton("Skip upload of ramdisk\n	(already used ramdisk method 1)")
		self.skip2 = QRadioButton("Skip upload of ramdisk\n	(already used ramdisk method 2)")
		vboxiThing.addWidget(self.threeGradio)
		vboxiThing.addWidget(self.threeGSradio)
		vboxiThing.addWidget(self.fourGSMradio)
		vboxiThing.addWidget(self.fourCDMAradio)
		vboxiThing.addWidget(self.pod2Gradio)
		vboxiThing.addWidget(self.podthreeGradio)
		vboxiThing.addWidget(self.pod4Gradio)
		vboxiThing.addWidget(self.pad1radio)
		vboxiThing.addWidget(self.skip1)
		vboxiThing.addWidget(self.skip2)
		continueiThingButton = QPushButton("Continue")
		vboxiThing.addWidget(continueiThingButton)
		self.connect(continueiThingButton, SIGNAL('clicked()'), self.setiThing)
		self.groupBoxiThing = QGroupBox("iThing")
		self.groupBoxiThing.setLayout(vboxiThing)
		self.stackedWidget.addWidget(self.groupBoxiThing)

		self.lineEditMan = QLineEdit()
		continueButtonMan = QPushButton("Continue")
		self.connect(continueButtonMan, SIGNAL('clicked()'), self.buildRamdisk)
		vboxMan = QVBoxLayout()
		vboxMan.addWidget(self.lineEditMan)
		vboxMan.addWidget(continueButtonMan)
		self.groupBoxMan = QGroupBox("Manual Path")
		self.groupBoxMan.setLayout(vboxMan)
		self.stackedWidget.addWidget(self.groupBoxMan)

		removeButton = QPushButton("Remove")
		skipButton = QPushButton("Skip")
		self.connect(removeButton, SIGNAL('clicked()'), self.removeIPSW)
		self.connect(skipButton, SIGNAL('clicked()'), self.skipRemoveIPSW)
		vboxRemove = QVBoxLayout()
		vboxRemove.addWidget(removeButton)
		vboxRemove.addWidget(skipButton)
		self.groupBoxRemove = QGroupBox("Remove?")
		self.groupBoxRemove.setLayout(vboxRemove)
		self.stackedWidget.addWidget(self.groupBoxRemove)

		yesRamdiskButton = QPushButton("Yes")
		noRamdiskButton = QPushButton("No")
		self.connect(yesRamdiskButton, SIGNAL('clicked()'), self.selRamdiskMeth)
		self.connect(noRamdiskButton, SIGNAL('clicked()'), self.manualPath)
		vboxRamdisk = QVBoxLayout()
		vboxRamdisk.addWidget(yesRamdiskButton)
		vboxRamdisk.addWidget(noRamdiskButton)
		self.groupBoxRamdisk = QGroupBox("Ramdisk Found")
		self.groupBoxRamdisk.setLayout(vboxRamdisk)
		self.stackedWidget.addWidget(self.groupBoxRamdisk)

		
		methodOneButton = QPushButton("Method One")
		methodTwoButton = QPushButton("Method Two")
		self.connect(methodOneButton, SIGNAL('clicked()'), lambda: self.redsn0w(1))
		self.connect(methodTwoButton, SIGNAL('clicked()'), lambda: self.redsn0w(2))
		vboxMeth = QVBoxLayout()
		vboxMeth.addWidget(methodOneButton)
		vboxMeth.addWidget(methodTwoButton)
		self.groupBoxMeth = QGroupBox("Ramdisk Method")
		self.groupBoxMeth.setLayout(vboxMeth)
		self.stackedWidget.addWidget(self.groupBoxMeth)

		self.crackRadio = QRadioButton("Crack four digit password")
		self.passRadio = QRadioButton("Get passwords")
		self.imageDataRadio = QRadioButton("Image data partition")
		self.imageSystemRadio = QRadioButton("Image system partition")
		self.recoverRadio = QRadioButton("Recover deleted files")
		self.continuePwnButton = QPushButton("Continue")
		vboxPwn1 = QVBoxLayout()
		vboxPwn1.addWidget(self.crackRadio)
		vboxPwn1.addWidget(self.passRadio)
		vboxPwn1.addWidget(self.imageDataRadio)
		vboxPwn1.addWidget(self.imageSystemRadio)
		vboxPwn1.addWidget(self.recoverRadio)
		vboxPwn1.addWidget(self.continuePwnButton)
		self.connect(self.continuePwnButton, SIGNAL('clicked()'), self.pwn)
		self.groupBoxPwn1 = QGroupBox("Method 1")
		self.groupBoxPwn1.setLayout(vboxPwn1)
		self.stackedWidget.addWidget(self.groupBoxPwn1)

		self.crack1Radio = QRadioButton("Crack four digit password\n	(can use if method 1 not working)")
		self.crack2Radio = QRadioButton("Crack four digit password\n	(can use if others not working)")
		self.shellRadio = QRadioButton("Drop to shell")
		self.picsRadio = QRadioButton("Get non-deleted pictures")
		self.dataFilesRadio = QRadioButton("Get all non-deleted files on data partition")
		self.systemFilesRadio = QRadioButton("Get all non-deleted files on system partition")
		self.rebootRadio = QRadioButton("Reboot")
		self.continuePwnButton2 = QPushButton("Continue")
		vboxPwn2 = QVBoxLayout()
		vboxPwn2.addWidget(self.crack1Radio)
		vboxPwn2.addWidget(self.crack2Radio)
		vboxPwn2.addWidget(self.shellRadio)
		vboxPwn2.addWidget(self.picsRadio)
		vboxPwn2.addWidget(self.dataFilesRadio)
		vboxPwn2.addWidget(self.systemFilesRadio)
		vboxPwn2.addWidget(self.continuePwnButton2)
		#self.connect(self.continuePwnButton2, SIGNAL('clicked()'), self.pwn)
		self.groupBoxPwn2 = QGroupBox("Method 2")
		self.groupBoxPwn2.setLayout(vboxPwn2)
		self.stackedWidget.addWidget(self.groupBoxPwn2)

#Connect

		


		



		mainHBox.addWidget(self.stackedWidget)
		self.stackedWidget.setContentsMargins(11, 11, 11, 11)
		mainWidget.setLayout(mainHBox)
		self.setCentralWidget(mainWidget)

		msg = "Click Setup if you have not before other wise Continue"
		self.printmsg(msg)

	def setup(self):

		for msg in dep.setup():
			self.printmsg(msg) #execute dependencies file and print returned messages

		try:
			check_call(['make', '-C', 'ramdisk_tools']) ####   <<------ Check

			self.printmsg(
'''------------------------------------------------------------------------------------------------------------------------
------------------------------------------------------------------------------------------------------------------------
Setup finished
------------------------------------------------------------------------------------------------------------------------
------------------------------------------------------------------------------------------------------------------------
'''
)
		except:
			self.printmsg(

'''
Failed making iphone-dataprotection/ramdisk_tools

You can try running:
make -C iphone-dataprotection/ramdisk_tools

in this directory in Terminal



cd [this directory]
to change to this directory

'''

)

	def displayiThing(self):
		# Shows set of iThing radio buttons
		# Continue ->  setiThing()
	
		self.stackedWidget.setCurrentWidget(self.groupBoxiThing)

		msg = \
'''

Select iThing.

'''
		self.printmsg(msg)

	def setiThing(self):
		# Set's iThing global
		# Usually -> chkRamdisk()

		global IPSW_name
		global ramdiskMeth
		print "hello from setiThing"
		
		if self.threeGradio.isChecked(): self.iThing = 'iPhone1,2'
                elif self.threeGSradio.isChecked(): self.iThing ='iPhone2,1'  
                elif self.fourGSMradio.isChecked(): self.iThing ='iPhone3,1'
                elif self.fourCDMAradio.isChecked(): self.iThing ='iPhone3,3'
                elif self.pod2Gradio.isChecked(): self.iThing ='iPod2,1'
                elif self.podthreeGradio.isChecked(): self.iThing ='iPod3,1'
                elif self.pod4Gradio.isChecked(): self.iThing ='iPod4,1'
                elif self.pad1radio.isChecked(): self.iThing ='iPad1,1'
		elif self.skip1.isChecked():
			self.iThing = 'skip'
			ramdiskMeth = '1'
		elif self.skip2.isChecked(): 
			oelf.iThing = 'skip'
			ramdiskMeth = '2'
		else:
			print "hello from else"



		if not self.iThing  == 'skip':  #Tired of skipping this part in testing
			IPSW_name = devices[self.iThing][1]	#set name of IPSW we are using
			self.chkRamdisk()
		else:
			self.redsn0w(ramdiskMeth)
		
	def chkRamdisk(self):
		global IPSW_path
		# If previous ramdisk found ask if we whant to use it.
		#	Yes -> selRamdiskMeth()				 
		#       No ->  manualPath()
		# Else manualPath()

		if exists('myramdisk_' + devices[self.iThing][3] + '.dmg'):

			IPSW_path = './IPSW'

			self.stackedWidget.setCurrentWidget(self.groupBoxRamdisk)

			msg = \
'''
Previous ramdisk found would you like to use it? 
'''
			self.printmsg(msg)
			#connect	build_ramdisk(IPSW_path)
		else:
			print "hello from exists"
			self.manualPath()  #Fuck the flow of this program is fucked, theres got to be a better way to do this

	def manualPath(self):

		self.stackedWidget.setCurrentWidget(self.groupBoxMan)
		msg = \
'''
If you would like to enter manual paths to search for IPSW's
enter them to the right with the path's seperrated by a ', '
(without quotes) then press continue.

'''
		self.printmsg(msg)
		
		

	def buildRamdisk(self):
		print "hello from build Ramdisk"
		global IPSW_files
		
		# Search for IPSW's put name and location in dictionary IPSW_files
		default_paths = [ getcwd() + "/IPSW" ]#   , expanduser("~/Library/iTunes/iPhone Software Updates") not working because of spaves, The first path will store hard links to itunes update folder as work around.
		itunes_ipsw_path = expanduser("~/Library/iTunes/iPhone Software Updates/")

		paths = self.lineEditMan.text()
		print "hello from line 313"
		
		if paths == '': 
			paths = default_paths
			for ipsw in listdir(default_paths[0]):      ## Refresh links in IPSW folder with ones from updates folder, adding the path to default_paths does not work because of spaces
					if islink(default_paths[0] + '/' + ipsw):
						remove(default_paths[0] + ipsw)

			for ipsw in listdir(itunes_ipsw_path):
				
				try:
					link(itunes_ipsw_path + ipsw, default_paths[0] + '/' + ipsw)
				except OSError:
					pass  #<--- because remove not working??
		else:
			paths = expanduser(paths).split(', ')
		print "hello from line 329"
		
		self.printmsg('\n\n\nOne Sec\n\n\n')

		IPSW_files = dict([(file, dir) for (files, dir) in ((filter(lambda x: '.ipsw' in x, listdir(path)), path) for path in paths) for file in files])# make a list containing (IPSW, LOCATION), then convert to dictionary, this could probably be done better

		self.printmsg(str(IPSW_files))
		self.printmsg("Done searching IPSW's")

#-----------------------------------------------
#Check for correct IPSW if not download it

		if devices[self.iThing][1] not in IPSW_files:
			print "Hello from line 339"
			self.printmsg( "Could not find appropriate IPSW. Downloading the " + devices[self.iThing][1] + " IPSW from Apple")
			download_IPSW(self.iThing)
			print "Hello from line 342"
		if (sha1(open(IPSW_files[devices[self.iThing][1]] + '/' + devices[self.iThing][1]).read()).hexdigest()) == devices[self.iThing][5]: # Uses 700 mb of memory... will fix later        
			IPSW_path = IPSW_files[IPSW_name]
			self.build_ramdisk(IPSW_path)
		else:
			
			self.stackedWidget.setCurrentWidget(self.groupBoxRemove)

			msg = \
'''
SHA1 did not check out, Downloading again. If you want to download it yourself just place it in the ./iphone-dataprotection/IPSW folder.

You may have to remove this file so we do not try to use it it next time, Do you want to do this now?  (y/n)

'''
			self.printmsg(msg)
			

	def removeIPSW(self):

		global IPSW_files
		remove(IPSW_files[devices[self.iThing][1]] + '/' + devices[self.iThing][1])
		self.printmsg( 'done')

		self.download_IPSW(self.iThing)
		self.build_ramdisk()
		self.stackedWidget.setCurrentWidget(self.groupBoxRamdisk)

	def skipRemoveIPSW(self):

		self.download_IPSW(self.iThing)
		self.build_ramdisk()
		self.stackedWidget.setCurrentWidget(self.groupBoxRamdisk)
		

	def selRamdiskMeth(self):

		self.stackedWidget.setCurrentWidget(self.groupBoxMeth)
		
		msg = \
'''
There is two different ways of uploading the ramdisk
Either one will allow you to do certain things:

Method 1
^^^^^^^^
1) Crack four digit password
2) Get passwords 
3) DD data Partition 
4) DD system Partition 
5) Recover Deleted files 

Method 2
^^^^^^^^
6) Crack four digit password (use if first one isn't working) 
7) Crack four digit password (if other two don't work) 
8) Get shell 
9) Get Pictures 
10) Get all non-deleted files on data partition
11) Get all non-deleted files on system partition
12) Get all non-deleted files
'''

		self.printmsg(msg)
		
	def redsn0w(self, meth):
		# Also Sets global ramdiskMeth

		global IPSW_path
		global ramdiskMeth

		print getcwd()
		print "hello from redsn0w " , meth
	
		self.printmsg( "\n\n Follow the directions in Redsn0w. When you see a big OK on your iDevice hit cancel in redsn0w\n\n")

		if meth == 1:
			ramdiskMeth = 1
			call(['./redsn0w_mac_0.9.12b2/redsn0w.app/Contents/MacOS/redsn0w', '-i', IPSW_path + '/' + devices[self.iThing][1], '-r', './myramdisk_' + devices[self.iThing][3] + '.dmg', '-k', devices[self.iThing][2],  '-a', '"-v rd=md0 nand-disable=1"'])
		if meth == 2:	
			ramdiskMeth = 2
			call(['./redsn0w_mac_0.9.12b2/redsn0w.app/Contents/MacOS/redsn0w', '-i',  IPSW_path + '/' + devices[self.iThing][1],  '-r', './myramdisk_' + devices[self.iThing][3] + '.dmg', '-k',  devices[self.iThing][2]])
			
		self.printmsg('Setting TCP relay')
		try:
			call(['python usbmuxd-python-client/tcprelay.py -t 22:2222 1999:1999 &'], shell = True)	
		except:
			print a
		self.printmsg('Set TCP relay')


		print meth
		
		if meth == 1:
			self.stackedWidget.setCurrentWidget(self.groupBoxPwn1)
		elif meth == 2:
			self.stackedWidget.setCurrentWidget(self.groupBoxPwn2)

		msg = \
'''

Select what you would like to do and
press enter.

'''
		self.printmsg(msg)
		
	def pwn(self):
		global ramdiskMeth
		
		if ramdiskMeth == 1:
			if self.crackRadio.isChecked():
				ios_examiner.shell.onecmd('bruteforce')
			elif self.passRadio.isChecked():
				ios_examiner.shell.onecmd('keychain')
			elif self.imageDataRadio.isChecked():
				ios_examiner.shell.onecmd('data')
				input = QInputDialog.getText(self, "File Name", "What do you want the file to be called", QLineEdit.Normal, '', 'ok')
				ios_examiner.shell.onecmd('dd ' + input)
			elif self.imageSystemRadio.isChecked():
				ios_examiner.shell.onecmd('system')
				input = QInputDialog.getText(self, "File Name", "What do you want the file to be called", QLineEdit.Normal, '', 'ok')
				ios_examiner.shell.onecmd('dd ' + input)
			elif self.recoverRadio.isChecked():
				ios_examiner.shell.onecmd('undelete')
			elif self.rebootRadio.isChecked():
				call('echo "reboot"|ssh -i ssh/id_rsa -p 2222 root@localhost', shell=True)
		elif ramdiskMeth == 2:
			if self.crack1Radio.isChecked():
				import demo_bruteforce
			elif self.crack2Radio.isChecked():
				getcwd()
				call("echo ./bruteforce | ssh -i ssh/id_rsa -p 2222 root@localhost", shell=True)
				ls = check_output("echo ls | ssh -i ssh/id_rsa -p 2222 root@localhost", shell=True)
				infos_file = search(r'[0-9, A-Z, a-z]*\.plist', ls).group(0)
				call("scp -i ssh/id_rsa -P 2222 root@localhost:" + infos_file + " .", shell=True)
				open(infos_file, 'r').read()
			elif self.shellRadio.isChecked():
				msg = \
'''
-------------------
-------------------
-------------------
Password is alpine
-------------------
-------------------
-------------------
'''

				self.printmsg(msg)
				call(["./iphone-dataprotection/ssh.command"], shell=True)
				
			elif self.picsRadio.isChecked():
				call('echo "./bruteforce;./mount_partitions.sh"|ssh -i ssh/id_rsa -p 2222 root@localhost', shell=True)
				call('scp -r -i ssh/id_rsa -P 2222 root@localhost:/mnt2/mobile/Media/DCIM/100APPLE .', shell=True)
				input = QInputDialog.getText(self, "File Name", "What do you want the file to be called", QLineEdit.Normal, '', 'ok')
				rename('100APPLE', input)
			elif self.dataFilesRadio.isChecked():
				call('echo "./bruteforce;./mount_partitions.sh"|ssh -i ssh/id_rsa -p 2222 root@localhost', shell=True)
				call('scp -r -i ssh/id_rsa -P 2222 root@localhost:/mnt2 .', shell=True)
				input = QInputDialog.getText(self, "File Name", "What do you want the file to be called", QLineEdit.Normal, '', 'ok')
				rename('mnt2', input('Name data partition folder'))
			elif self.systemFilesRadio.isChecked():
				call('echo "./bruteforce;./mount_partitions.sh"|ssh -i ssh/id_rsa -p 2222 root@localhost', shell=True)
				call('scp -r -i ssh/id_rsa -P 2222 root@localhost:/mnt1 .', shell=True)
				input = QInputDialog.getText(self, "File Name", "What do you want the file to be called", QLineEdit.Normal, '', 'ok')
				rename('mnt1', input)
			elif self.rebootRadio.isChecked():
				call('echo "reboot"|ssh -i ssh/id_rsa -p 2222 root@localhost', shell=True)
	def printmsg(self, msg):
		self.listWidget.addItem(msg)
		self.listWidget.scrollToBottom()

	def build_ramdisk(self, IPSW_path = './IPSW'):
		kernel_patcher.main(IPSW_path + '/' + devices[self.iThing][1], '')
		call(['sh ./make_ramdisk_' + devices[self.iThing][3] + '.sh'], shell = True)

	def download_IPSW(self, iThing):
		print "hello from download_IPSW"
		urlretrieve(devices[iThing][4], "./IPSW/" + devices[iThing][1], reporthook=self.dlProgress)

	def dlProgress(self, count, blockSize, totalSize):
		percent = int(count*blockSize*100/totalSize)
		print "\r%2d%%" % percent


def main():
	app = QApplication(argv)
	app.setApplicationName("Data Recovery")
	form = MainWindow()
	form.show()
	app.exec_()




main()
